import os
import sys
import time
from kubernetes import client, config
from kubernetes.client.rest import ApiException

config.load_incluster_config()
batch_v1 = client.BatchV1Api()
core_v1  = client.CoreV1Api()

def get_node_info():
  node_info = {}
  nodes = core_v1.list_node()

  for node in nodes.items:
    node_name = node.metadata.name
    is_master = "node-role.kubernetes.io/master" in node.metadata.labels
    node_ip = None
    for address in node.status.addresses:
      if address.type == "InternalIP":
        node_ip = address.address
        break
  
    allocatable = node.status.allocatable
    capacity = node.status.capacity
  
    node_info[node_ip] = {
      'is_master': is_master,
      'hostname': node_name,
      'ip': node_ip,
      'cpu': {
        'allocatable': allocatable.get('cpu'),
        'capacity': capacity.get('cpu')
      },
      'memory': {
        'allocatable': allocatable.get('memory'),
        'capacity': capacity.get('memory')
      },
      'ephemeral_storage': {
        'allocatable': allocatable.get('ephemeral-storage'),
        'capacity': capacity.get('ephemeral-storage')
      },
      'hugepages_1Gi': {
        'allocatable': allocatable.get('hugepages-1Gi'),
        'capacity': capacity.get('hugepages-1Gi')
      },
      'hugepages_2Mi': {
        'allocatable': allocatable.get('hugepages-2Mi'),
        'capacity': capacity.get('hugepages-2Mi')
      },
      'pods': {
        'allocatable': allocatable.get('pods'),
        'capacity': capacity.get('pods')
      }
    }

  return node_info

